import cv2
import numpy as np
import os

class MangaPanelExtractor:
    def __init__(self, min_panel_size=1000):
        self.min_panel_size = min_panel_size
        # 设置分割线的最小长宽比
        self.min_aspect_ratio = 3.0  # 分割线的长度至少是宽度的3倍
        # 设置分割线的最小长度（相对于图片宽度的比例）
        self.min_line_length_ratio = 0.6  # 分割线长度至少是图片宽度的60%

    def preprocess_image(self, image):
        # 转换为灰度图
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        
        # 使用自适应阈值处理，这样可以更好地处理不同亮度区域
        binary = cv2.adaptiveThreshold(
            gray,
            255,
            cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
            cv2.THRESH_BINARY,
            11,  # 邻域大小
            2    # 常数差值
        )
        
        # 使用形态学操作来增强水平线条
        # 创建水平结构元素
        horizontal_size = int(binary.shape[1] / 30)
        horizontal_structure = cv2.getStructuringElement(cv2.MORPH_RECT, (horizontal_size, 1))
        
        # 腐蚀和膨胀操作，增强水平线条
        eroded = cv2.erode(binary, horizontal_structure)
        dilated = cv2.dilate(eroded, horizontal_structure)
        
        return dilated

    def filter_components(self, binary_image):
        # 标记连通区域
        num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(binary_image, connectivity=8)
        print(f"找到 {num_labels-1} 个连通区域")
        
        height, width = binary_image.shape
        min_line_length = int(width * self.min_line_length_ratio)
        print(f"最小线条长度: {min_line_length} 像素")
        
        # 创建掩码图像
        mask = np.zeros_like(binary_image)
        
        # 分析每个连通区域
        separator_count = 0
        for i in range(1, num_labels):  # 跳过背景（标签0）
            x, y, w, h, area = stats[i]
            
            # 计算长宽比
            aspect_ratio = max(w, h) / min(w, h) if min(w, h) > 0 else 0
            
            # 判断是否是分割线
            is_separator = (
                aspect_ratio >= self.min_aspect_ratio and  # 长宽比满足要求
                w >= min_line_length and                   # 宽度足够长
                h <= height * 0.1                          # 高度不能太大
            )
            
            if is_separator:
                separator_count += 1
                print(f"找到分割线 {separator_count}: 位置=({x}, {y}), 宽度={w}, 高度={h}, 长宽比={aspect_ratio:.2f}")
                # 将符合条件的区域添加到掩码中
                mask[labels == i] = 255
        
        print(f"总共找到 {separator_count} 条有效分割线")
        return mask

    def detect_panels(self, preprocessed_image):
        height, width = preprocessed_image.shape[:2]
        print(f"\n开始检测分镜，图片尺寸: {width}x{height}")
        
        # 过滤连通区域，只保留可能的分割线
        filtered_image = self.filter_components(preprocessed_image)
        
        # 计算水平方向的白色像素投影
        horizontal_proj = np.sum(filtered_image, axis=1) / 255
        
        # 找到主要的分隔位置
        threshold = width * 0.6  # 降低阈值到60%，提高灵敏度
        print(f"分割线检测阈值: {threshold} 像素")
        
        split_positions = []
        i = 0
        while i < height:
            if horizontal_proj[i] > threshold:
                # 找到分隔区域的起始位置
                start = i
                while i < height and horizontal_proj[i] > threshold:
                    i += 1
                # 分隔区域的中心位置
                center = (start + i) // 2
                split_positions.append(center)
                print(f"找到分割位置: y={center}")
            i += 1
        
        print(f"\n初始分割位置数量: {len(split_positions)}")
        
        # 合并太近的分割线
        min_distance = height * 0.05  # 最小间距为图片高度的5%
        filtered_positions = []
        for pos in split_positions:
            if not filtered_positions or abs(pos - filtered_positions[-1]) >= min_distance:
                filtered_positions.append(pos)
        
        print(f"合并后的分割位置数量: {len(filtered_positions)}")
        
        # 根据分隔位置创建面板
        valid_panels = []
        if filtered_positions:
            prev_pos = 0
            for pos in filtered_positions:
                panel_height = pos - prev_pos
                if panel_height >= self.min_panel_size:  # 确保面板大小足够
                    valid_panels.append((0, prev_pos, width, panel_height))
                    print(f"添加面板: y={prev_pos}, 高度={panel_height}")
                else:
                    print(f"跳过过小的面板: y={prev_pos}, 高度={panel_height}")
                prev_pos = pos
            # 添加最后一个面板
            final_height = height - prev_pos
            if final_height >= self.min_panel_size:
                valid_panels.append((0, prev_pos, width, final_height))
                print(f"添加最后一个面板: y={prev_pos}, 高度={final_height}")
            else:
                print(f"跳过过小的最后一个面板: y={prev_pos}, 高度={final_height}")
        else:
            # 如果没有找到分隔线，保持整页
            valid_panels.append((0, 0, width, height))
            print("没有找到分隔线，保持整页")
        
        print(f"\n最终检测到 {len(valid_panels)} 个有效面板")
        return valid_panels

    def extract_panels(self, image_path, output_dir):
        print(f"\n开始处理图片: {image_path}")
        
        # 创建输出目录
        os.makedirs(output_dir, exist_ok=True)
        print(f"输出目录: {output_dir}")
        
        # 读取图片
        image = cv2.imread(image_path)
        if image is None:
            raise ValueError(f"Could not read image: {image_path}")
        
        # 获取图片尺寸
        height, width = image.shape[:2]
        print(f"图片尺寸: {width}x{height}")
        
        # 预处理图片
        print("\n开始预处理图片...")
        preprocessed = self.preprocess_image(image)
        
        # 保存预处理后的图片用于调试
        preprocessed_path = os.path.join(output_dir, "preprocessed.jpg")
        cv2.imwrite(preprocessed_path, preprocessed)
        print(f"保存预处理图片: {preprocessed_path}")
        
        # 保存过滤后的图片用于调试
        print("\n开始过滤连通区域...")
        filtered = self.filter_components(preprocessed)
        filtered_path = os.path.join(output_dir, "filtered.jpg")
        cv2.imwrite(filtered_path, filtered)
        print(f"保存过滤后图片: {filtered_path}")
        
        # 检测分镜
        panels = self.detect_panels(preprocessed)
        
        # 在原图上绘制检测到的分镜位置
        debug_image = image.copy()
        for i, (x, y, w, h) in enumerate(panels):
            cv2.rectangle(debug_image, (x, y), (x+w, y+h), (0, 255, 0), 2)
            cv2.putText(debug_image, str(i+1), (x+10, y+30), 
                       cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
        
        # 保存标注后的图片
        debug_path = os.path.join(output_dir, "detected_panels.jpg")
        cv2.imwrite(debug_path, debug_image)
        print(f"保存调试图片: {debug_path}")
        
        # 提取并保存每个分镜
        print("\n开始保存分镜...")
        for i, (x, y, w, h) in enumerate(panels):
            # 提取分镜
            panel = image[y:y+h, x:x+w]
            
            # 保存分镜
            output_path = os.path.join(output_dir, f"panel_{i+1}.jpg")
            cv2.imwrite(output_path, panel)
            print(f"保存分镜 {i+1}: {output_path}")
        
        return len(panels)

def main():
    # 创建 MangaPanelExtractor 实例，将最小分镜大小设置为较小的值
    extractor = MangaPanelExtractor(min_panel_size=500)  # 改为500像素
    
    # 确保输入和输出目录存在
    input_dir = "pic"
    output_base_dir = "extracted_panels"
    
    if not os.path.exists(input_dir):
        os.makedirs(input_dir)
        print(f"创建输入目录: {input_dir}")
    
    if not os.path.exists(output_base_dir):
        os.makedirs(output_base_dir)
        print(f"创建输出目录: {output_base_dir}")
    
    # 获取所有 PNG 文件
    png_files = [f for f in os.listdir(input_dir) if f.lower().endswith('.png') or f.lower().endswith('.jpg')]
    print(f"\n找到 {len(png_files)} 个PNG文件")
    
    if not png_files:
        print(f"\n错误：在 {input_dir} 目录中没有找到PNG文件！")
        print("请确保：")
        print("1. 图片文件已经放在 pic 目录中")
        print("2. 图片文件是PNG格式的")
        print("3. 文件扩展名是.png（小写）")
        return
    
    # 处理每个 PNG 文件
    for png_file in png_files:
        print(f"\n处理文件: {png_file}")
        input_path = os.path.join(input_dir, png_file)
        
        # 创建以 PNG 文件名（不含扩展名）命名的输出目录
        output_dir = os.path.join(output_base_dir, os.path.splitext(png_file)[0])
        
        try:
            # 如果输出目录已存在，先删除
            if os.path.exists(output_dir):
                import shutil
                shutil.rmtree(output_dir)
                print(f"清理已存在的输出目录: {output_dir}")
            
            # 处理图片
            num_panels = extractor.extract_panels(input_path, output_dir)
            print(f"成功从 {png_file} 中提取了 {num_panels} 个分镜到 {output_dir}")
            
        except Exception as e:
            print(f"处理 {png_file} 时出错: {str(e)}")
            continue

if __name__ == "__main__":
    main() 